/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2024-2024.
 * Description: rtos kexec hardening for security
 * Author: chenjingwen
 * Create: 2024-02-02
 */
#include <crypto/sha2.h>
#include <crypto/hash.h>
#include <linux/string.h>
#include <linux/printk.h>
#include <linux/libfdt.h>
#include <linux/of_fdt.h>
#include <linux/init.h>
#include <internal/rtos_sec_internal.h>

#define DIGEST_SIZE SHA256_DIGEST_SIZE

static u8 dtb_hash[DIGEST_SIZE];

int check_dtb_integrity(const void *dtb_addr, unsigned int dtb_size)
{
	u8 tmp_dtb_hash[DIGEST_SIZE] = { 0 };
	int ret;

	if (!dtb_addr || dtb_size == 0)
		return -EINVAL;

	ret = rtos_calc_sha256((const u8 *)dtb_addr, dtb_size, tmp_dtb_hash);
	if (ret) {
		pr_err("[kexec_hardening]: cannot calculate hash for dtb\n");
		return ret;
	}

	if (memcmp(dtb_hash, tmp_dtb_hash, DIGEST_SIZE) != 0) {
		pr_err("[kexec_hardening]: The given dtb has been modified\n");
		return -EPERM;
	}

	return 0;
}

static int __init store_dtb_integrity(const void *dtb_addr,
				      unsigned int dtb_size)
{
	if (!dtb_addr || dtb_size == 0)
		return -EINVAL;

	return rtos_calc_sha256((const u8 *)dtb_addr, dtb_size, dtb_hash);
}

static int __init kexec_hardening_init(void)
{
	int ret = store_dtb_integrity(initial_boot_params,
				      fdt_totalsize(initial_boot_params));
	if (ret)
		pr_err("[kexec_hardening]: unable to store dtb hash, r:%d\n", ret);

	return ret;
}
late_initcall(kexec_hardening_init);
